#include <stdio.h>
#include <string.h>
#include <time.h>
#include <stdlib.h>

#include "Algorithm\RSA2048\rsa.h"
#include "Algorithm\RSA2048\keys.h"
#include "Algorithm\Convert\ConvertDataType.h"

#define DATA_TYPE_CONVERT (1)
#define READ_FILE_CALC_RSA (1)
#define SHA256_ENABLE (1)

/*
 * RSA2048 encrypt and decrypt
 * include rsa.c/bignum.c/rsa.h/bignum.h/keys.h
 */
static int RSA2048(void){
    int ret;
    rsa_pk_t pk = {0};
    rsa_sk_t sk = {0};
    uint8_t output[256];

    // message to encrypt
    uint8_t input [256] = { 0x21,0x55,0x53,0x53,0x53,0x53};

    unsigned char msg [256];
    uint32_t outputLen, msg_len;
    uint8_t  inputLen;

    // copy keys.h message about public key and private key to the flash RAM
    pk.bits = KEY_M_BITS;
    memcpy(&pk.modulus         [RSA_MAX_MODULUS_LEN-sizeof(key_m) ],  key_m,  sizeof(key_m ));
    memcpy(&pk.exponent        [RSA_MAX_MODULUS_LEN-sizeof(key_e) ],  key_e,  sizeof(key_e ));
    sk.bits = KEY_M_BITS;
    memcpy(&sk.modulus         [RSA_MAX_MODULUS_LEN-sizeof(key_m) ],  key_m,  sizeof(key_m ));
    memcpy(&sk.public_exponet  [RSA_MAX_MODULUS_LEN-sizeof(key_e) ],  key_e,  sizeof(key_e ));
    memcpy(&sk.exponent        [RSA_MAX_MODULUS_LEN-sizeof(key_pe)],  key_pe, sizeof(key_pe));
    memcpy(&sk.prime1          [RSA_MAX_PRIME_LEN - sizeof(key_p1)],  key_p1, sizeof(key_p1));
    memcpy(&sk.prime2          [RSA_MAX_PRIME_LEN - sizeof(key_p2)],  key_p2, sizeof(key_p2));
    memcpy(&sk.prime_exponent1 [RSA_MAX_PRIME_LEN - sizeof(key_e1)],  key_e1, sizeof(key_e1));
    memcpy(&sk.prime_exponent2 [RSA_MAX_PRIME_LEN - sizeof(key_e2)],  key_e2, sizeof(key_e2));
    memcpy(&sk.coefficient     [RSA_MAX_PRIME_LEN - sizeof(key_c) ],  key_c,  sizeof(key_c ));

    inputLen = strlen((const char*)input);

    // public key encrypt
    rsa_public_encrypt(output, &outputLen, input, inputLen, &pk);
    printf("public key encrypt is %d\n",outputLen);
    TRACE(output,outputLen);

    // private key decrypt
    rsa_private_decrypt(msg, &msg_len, output, outputLen, &sk);
    printf("private key decrypt  len is %d\n",msg_len);
    TRACE(msg,msg_len);

    // private key encrypt
    rsa_private_encrypt(output, &outputLen, input, inputLen, &sk);
    printf("private key encrypt len is %d\n",outputLen);
    TRACE(output,outputLen);

    // public key decrypted
    rsa_public_decrypt(msg, &msg_len, output, outputLen, &pk);
    printf("public key decrypted len is %d\n",msg_len);
    TRACE(msg,msg_len);

    return 0;
}
/* RSA2048 function ended */

// message to encrypt
uint8_t inputkey [32] = { 0xa1, 0xfe, 0x12, 0x73, 0x33, 0x16, 0xc7, 0xc7, 0x02, 0xed, 0xcf, 0x00, 0xb4, 0x5b, 0xde, 0xbe }; //AES128 EnCode Key
uint32_t inputkey_Len = 32; //input key len
uint8_t outputkey[256];
uint32_t outputkey_Len = 256;
uint8_t msgkey [32];
uint32_t msgkey_Len = 32;

/**
  * @brief  RSA2048加密
  * @param  plaintext 明文
  * @param  plaintext_Len 明文长度
  * @param  ciphertext 密文
  * @param  ciphertext_Len 密文长度
  * @return none
  * @note   none
  */
void RSA2048EnCode(uint8_t plaintext[256], uint32_t plaintext_Len, uint8_t ciphertext[256], uint32_t ciphertext_Len)
{
    rsa_sk_t sk = {0};

    // copy keys.h message about public key and private key to the flash RAM
    sk.bits = KEY_M_BITS;
    memcpy(&sk.modulus         [RSA_MAX_MODULUS_LEN-sizeof(key_m) ],  key_m,  sizeof(key_m ));
    memcpy(&sk.public_exponet  [RSA_MAX_MODULUS_LEN-sizeof(key_e) ],  key_e,  sizeof(key_e ));
    memcpy(&sk.exponent        [RSA_MAX_MODULUS_LEN-sizeof(key_pe)],  key_pe, sizeof(key_pe));
    memcpy(&sk.prime1          [RSA_MAX_PRIME_LEN - sizeof(key_p1)],  key_p1, sizeof(key_p1));
    memcpy(&sk.prime2          [RSA_MAX_PRIME_LEN - sizeof(key_p2)],  key_p2, sizeof(key_p2));
    memcpy(&sk.prime_exponent1 [RSA_MAX_PRIME_LEN - sizeof(key_e1)],  key_e1, sizeof(key_e1));
    memcpy(&sk.prime_exponent2 [RSA_MAX_PRIME_LEN - sizeof(key_e2)],  key_e2, sizeof(key_e2));
    memcpy(&sk.coefficient     [RSA_MAX_PRIME_LEN - sizeof(key_c) ],  key_c,  sizeof(key_c ));

    //plaintext
    // printf("plaintext is:");
    // TRACE(plaintext,plaintext_Len);

    // private key encrypt
    rsa_private_encrypt(ciphertext, &ciphertext_Len, plaintext, plaintext_Len, &sk);
    printf("encrypt data is %d:",ciphertext_Len);
    TRACE(ciphertext,ciphertext_Len);
}

/**
  * @brief  RSA2048解密
  * @param  ciphertext 密文
  * @param  ciphertext_Len 密文长度
  * @param  plaintext 明文
  * @param  plaintext_Len 明文长度
  * @return none
  * @note   none
  */
void RSA2048DeCode(uint8_t ciphertext[256], uint32_t ciphertext_Len, uint8_t plaintext[32], uint32_t plaintext_Len)
{
    rsa_pk_t pk = {0};

    // copy keys.h message about public key and private key to the flash RAM
    pk.bits = KEY_M_BITS;
    memcpy(&pk.modulus         [RSA_MAX_MODULUS_LEN-sizeof(key_m) ],  key_m,  sizeof(key_m ));
    memcpy(&pk.exponent        [RSA_MAX_MODULUS_LEN-sizeof(key_e) ],  key_e,  sizeof(key_e ));

    // public key decrypt
    rsa_public_decrypt(plaintext, &plaintext_Len, ciphertext, ciphertext_Len, &pk);
    printf("public key decrypt  len is %d\n",plaintext_Len);
    TRACE(plaintext,plaintext_Len);
}

#if (READ_FILE_CALC_RSA)

#if (SHA256_ENABLE)
#include "Algorithm\SHA256\Sha_256.h"
STATIC Sha256Calc Sha256CalcValue;

void SecM_Sha256CalcInit(void)
{
  (void)Sha256Calc_init(&Sha256CalcValue);
}

uint8 *SecM_Sha256CalcCalculate(const uint8 *dp, uint32 dl)
{
  Sha256Calc_calculate(&Sha256CalcValue, dp, dl);
  return Sha256CalcValue.Value;
}
#endif

void remove0xAndCommas(char *str,char *des) {
    int i, j;
    for (i = 0, j = 0; str[i]; i++) {
        if (str[i] == '0' && str[i+1] == 'x')
        {
            i++;
        }
        else if (str[i] == ' ' || str[i] == ',' || str[i] == '	')
        {

        }
        else
        {
            des[j]=str[i];
            j++;
        }
    }
    // str[j] = '\0';
}

#define MAX_LINE_LENGTH 128
int ReadFileHashData(char *filename, char *targetArrayName, char *readhashdata)
{
  FILE *cFile;
  // char filename[] = "flashdriver.c";  // 用您的C文件的实际文件名替换此处
  // char targetArrayName[] = "Blk0";  // 替换为您的目标数组名称
  char *arrayData[3] = {NULL, NULL, NULL};
  int lineCount = 0;

  cFile = fopen(filename, "r");
  if (cFile == NULL) {
      // printf("could not open file %s\n", filename);
      return 1;
  }

  char line[MAX_LINE_LENGTH];
  char *arrayStart = NULL;
  int insideArray = 0;

  while (fgets(line, sizeof(line), cFile)) {
      // 在数组内
      if (insideArray) {
          if (arrayData[lineCount % 3] != NULL) {
              free(arrayData[lineCount % 3]);
          }
          arrayData[lineCount % 3] = strdup(line);
          lineCount++;
      } else {
          // 检查数组开始
          if ((arrayStart = strstr(line, targetArrayName)) != NULL) {
              insideArray = 1;
          }
      }
  }

  fclose(cFile);

  if (lineCount >= 3) {
      // printf("array end 3:\n%s", arrayData[(lineCount - 3) % 3]);
  }

  if (lineCount >= 2) {
      // printf("array end 2:\n%s", arrayData[(lineCount - 2) % 3]);
  }

  // 释放动态分配的内存
  for (int i = 0; i < 3; i++) {
      if (arrayData[i] != NULL) {
          free(arrayData[i]);
      }
  }

  char des1[33]={0};
  char des2[33]={0};
  remove0xAndCommas(arrayData[(lineCount - 3) % 3],des1);
  remove0xAndCommas(arrayData[(lineCount - 2) % 3],des2);

  // get file hash data
  for (int i = 0; i < 64; i++)
  {
    if (i<32)
    {
      readhashdata[i]=des1[i];
    }
    else
    {
      readhashdata[i]=des2[i-32];
    }
  }

  return 0;
}

void writeHexRecord(FILE *hexFile, uint32_t address, uint8_t *data, int dataSize) {
    fprintf(hexFile, ":");
    fprintf(hexFile, "%02X", dataSize);
    fprintf(hexFile, "%04X", address);
    fprintf(hexFile, "00");  // Record Type (data record)

    for (int i = 0; i < dataSize; i++) {
        fprintf(hexFile, "%02X", data[i]);
    }

    uint8_t checksum = dataSize;
    checksum += (address & 0xFF) + ((address >> 8) & 0xFF);
    for (int i = 0; i < dataSize; i++) {
        checksum += data[i];
    }
    checksum = (uint8_t)(0x100 - checksum);
    fprintf(hexFile, "%02X\n", checksum);
}

int main(int argc, char *argv[])
{
  int ret = 1;
  if (argc != 3)
  {
    printf("Description:\n");
    printf("  Note: Encryption - Hex + PKCS1 + RSA2048\n");
    printf("  1.Please store the hash value at the end of hex and convert the hex file to a C-Array file.\n");
    printf("  2.The program will read the last two rows of data (hash values 32Bytes) from the specified array in the specified file.\n");
    printf("  3.Create Output.hex to store encrypted data (256Bytes).\n");
    printf("  4.Online encrypted website URL:https://the-x.cn/cryptography/Rsa.aspx .\n");
    printf("Program Usage:\n");
    printf("  SHA256_To_RSA2048_Tool.exe Filename TargetArrayName\n");
    ret = 1;
  }
  else
  {
    printf("Read File: %s\n",argv[1]);
    printf("Read Array: %s[] (Get the last 2 rows hash data of array)\n\n",argv[2]);

    // Get file hash data
    char readhashdata[64]={0};
    char array_sha256[SHA256_LENGTH]={0};
    ReadFileHashData(argv[1],argv[2],readhashdata);
    string2hex(readhashdata, array_sha256);
    printf("SHA256 Data is:");
    TRACE(array_sha256,SHA256_LENGTH);

    // Calc rsa2048
    uint8_t encryptdata[256];
    uint32_t encryptdata_Len = 256;
    RSA2048EnCode(array_sha256,SHA256_LENGTH,encryptdata,encryptdata_Len);

    // Create Hex File
    FILE *hexFile;
    char filename[] = "Output.hex";  // 要创建的HEX文件的名称

    hexFile = fopen(filename, "w");
    if (hexFile == NULL) {
        printf("Create hex file %s fail\n", filename);
        return 1;
    }

    // 指定要存储数据的地址
    uint32_t address = 0x0000;
    for (int i = 0; i < encryptdata_Len/16; i++)
    {
        writeHexRecord(hexFile, address, encryptdata+(i*16), 16);
        address+=16;
    }

    fclose(hexFile);
    printf("Create hex file %s success\n", filename);

    ret = 0;
  }

  return ret;
}

#else
int main(int argc, char const *argv[])
{
    if (DATA_TYPE_CONVERT)
    {
        /* signture */
        char *input_sha256 = "7D22EB415CD5CB8D7F15BF51DA0BDA33F9A1B534E7B6D58F24F12FA9D3028A20";
        char array_sha256[32]={0};
        string2hex(input_sha256, array_sha256);
        memcpy(inputkey,array_sha256,32);
        printf("input_sha256:");
        TRACE(inputkey,32);
    }
    
    clock_t start, finish;
    double  duration;
    start = clock();    // init start time
    if (0)
    {
        RSA2048(); //test running   
    }
    else
    {
        RSA2048EnCode(inputkey,inputkey_Len,outputkey,outputkey_Len);
        RSA2048DeCode(outputkey,outputkey_Len,msgkey,msgkey_Len);
        TRACE(msgkey,32);
    }

    finish = clock();   // print end time
    duration = (double)(finish - start) / CLOCKS_PER_SEC;   // print encrypt and decrypt time
    printf( "%f seconds\n", duration );
    return 0;
}
#endif
